import random
from collections import deque

import torch

from light_rl.utils.utils import load_config, get_device
import os
import pickle
from typing import Union, Tuple
import numpy as np


class ReplayBuffer:
    def __init__(self, device_: torch.device, capacity: int):
        """
        经验池
        :param device:  torch.device 设备
        :param capacity:  int 容量
        """
        self.device = device_
        self.capacity = capacity
        self.buffer = deque(maxlen=self.capacity)

    def push(self, *data: Union[np.ndarray, int, bool]) -> None:
        """
        存放数据
        :param data: tuple 这里没有指定data的类型，可以自己定义
        :return:  None
        """
        self.buffer.append(data)

    def sample(self, batch_size: int) -> Tuple[torch.Tensor, ...]:
        """
        采样
        :param batch_size: int 采样的大小
        :return: zip 采样的数据
        """
        batch_indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        batch = [self.buffer[i] for i in batch_indices]
        sample_data = zip(*batch)
        tensor_data = ()
        for data in sample_data:
            item = torch.tensor(np.array(data), device=self.device)
            tensor_data += (item,)
        return tensor_data

    def __len__(self):
        return len(self.buffer)

    def save(self, path: str) -> None:
        """
        保存数据
        :param path: str 保存的路径
        :return: None
        """
        dir_name = os.path.dirname(path)
        if not os.path.exists(dir_name):
            os.mkdir(dir_name)
        with open(path, 'wb') as file:
            pickle.dump(self.buffer, file)

    def load(self, path: str) -> bool:
        """
        加载数据
        :param path: str 加载的路径
        :return: bool 是否加载成功
        """
        if os.path.exists(path):
            with open(path, 'rb') as file:
                self.buffer = pickle.load(file)
            return True
        else:
            return False


class EpochBuffer:
    def __init__(self, device: torch.device):
        self.buffer = []
        self.device = device

    def push(self, *data: Union[np.ndarray, int, bool]) -> None:
        """
        存放数据
        :param data: Tuple 这里没有指定data的类型，可以自己定义
        :return:  None
        """
        self.buffer.append(data)

    def get(self) -> Tuple[torch.Tensor, ...]:
        """
        获取数据
        :return: zip 采样的数据
        """
        zip_data = zip(*self.buffer)
        tensor_data = ()
        for data in zip_data:
            item = torch.tensor(np.array(data), device=self.device)
            tensor_data += (item,)
        return tensor_data

    def clear(self):
        self.buffer = []

    def __len__(self):
        return len(self.buffer)


def create_reply_buffer(config_file):
    config = load_config(config_file)
    capacity = config.buffer.buffer_size
    device_ = get_device(config.device)
    return ReplayBuffer(device_, capacity)


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    buffer = EpochBuffer(device)
    state = np.array([1, 2, 3])
    action = np.array(1)
    reward = 1
    next_state = np.array([4, 5, 6])
    done = False
    buffer.push(state, action, reward, next_state, done)
    buffer.push(state, action, reward, next_state, done)
    buffer.push(state, action, reward, next_state, done)

    states, actions, rewards, next_states, dones = buffer.get()
    pass
    # print(buffer.get())
